# youtube_subtitle_fetcher.py
import os
import requests
from bs4 import BeautifulSoup
import logging
import xml.etree.ElementTree as ET

class YouTubeSubtitleFetcher:
    def __init__(self, output_dir="output/subtitles"):
        """
        初始化字幕下载器
        :param output_dir: 输出目录
        """
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)

    def get_available_languages(self, video_id):
        """
        获取视频支持的字幕语言列表
        :param video_id: YouTube 视频 ID
        :return: 语言代码列表（如 ['en', 'zh-CN']）
        """
        # 尝试获取字幕列表的URL
        url = f"https://www.youtube.com/api/timedtext?v={video_id}&type=list"
        
        try:
            response = requests.get(url)
            if response.status_code == 200:
                # 解析XML格式的字幕列表
                root = ET.fromstring(response.text)
                lang_list = []
                for transcript in root.findall('transcript'):
                    lang_code = transcript.get('lang_code')
                    if lang_code:
                        lang_list.append(lang_code)
                
                if lang_list:
                    logging.info(f"🌐 视频 {video_id} 支持的语言: {lang_list}")
                    return lang_list
        except Exception as e:
            logging.debug(f"获取字幕列表失败: {e}")

        logging.warning(f"⚠️ 无法获取视频 {video_id} 的字幕列表")
        return []

    def download_subtitles(self, video_id, lang="en"):
        """
        下载 YouTube 字幕
        :param video_id: YouTube 视频 ID
        :param lang: 字幕语言，默认 en
        :return: 字幕文件路径
        """
        # 使用测试中发现的有效URL格式
        url = f"https://www.youtube.com/api/timedtext?v={video_id}&lang={lang}"
        
        try:
            response = requests.get(url)
            if response.status_code == 200 and response.text.strip():
                output_path = os.path.join(self.output_dir, f"{video_id}_{lang}.srt")
                
                with open(output_path, "w", encoding="utf-8") as f:
                    f.write(response.text)

                logging.info(f"📄 {lang}字幕已保存: {output_path}")
                return output_path
            else:
                logging.debug(f"获取字幕失败，状态码: {response.status_code}")
        except Exception as e:
            logging.error(f"获取字幕时出错: {e}")

        logging.warning(f"⚠️ 无法获取字幕: {video_id}")
        return None